package fem2.strategies;

import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;
import inf.jlinalg.SolveFailedException;

import java.util.Arrays;

import math2.MathUtilities;
import fem2.Debugger;
import fem2.Model;
import fem2.enu.EchoLevelType;
import fem2.enu.MassMatrixType;
import fem2.jlinalg.Solver;

public abstract class ImplicitTimeIntegrationStrategy extends Strategy {

	protected double alpha1 = 0.0;
	protected double alpha2 = 0.0;

	protected MassMatrixType massType;

	protected ArrayVector uHat = null;
	protected ArrayVector udHat = null;
	protected ArrayVector uddHat = null;

	protected ArrayVector uHat_old = null;
	protected ArrayVector udHat_old = null;
	protected ArrayVector uddHat_old = null;

	protected double t; // current time variable
	protected double t_old; // history time variable

	/**
	 * 
	 * @param m
	 * @param alpha1
	 *            Rayleigh damping coefficient
	 * @param alpha2
	 *            Rayleigh damping coefficient
	 */
	public ImplicitTimeIntegrationStrategy(Model m, MassMatrixType massType, double alpha1,
			double alpha2) {
		super(m);
		this.alpha1 = alpha1;
		this.alpha2 = alpha2;
		this.massType = massType;
	}

	@Override
	public double[] getCurrentUHat() {
		return uHat.getValues();
	}

	@Override
	public double[] getCurrentUdHat() {
		return udHat.getValues();
	}

	@Override
	public double[] getCurrentUddHat() {
		return uddHat.getValues();
	}

	/**
	 * initialize dependent matrices and vectors for specific dynamics
	 * integration schem
	 * 
	 * @param leftHandSideSolver
	 * @param rightHandSide
	 * @param u
	 */
	public abstract void initializeSystemMatricesAndVectors(Solver leftHandSideSolver,
			ArrayVector rightHandSide, ArrayVector u);

	/**
	 * 
	 * @return the next lambda for time step iteration
	 */
	public abstract double getNextTimeStep();

	/**
	 * compute initial condition and assign solution also finalize material
	 * parameters for the first step<br>
	 * compute uHat, udHat, uddHat<br>
	 * this is very important step. Please be careful
	 */
	public abstract void computeInitialCondition();

	/**
	 * assemble effective stiffness matrix and effective load vector
	 * 
	 * @param kteff
	 * @param reff
	 */
	public abstract void assembleEffectiveSystem(IMatrix kteff, IVector reff);

	/**
	 * update solution vector for the strategy. Different time integration
	 * scheme (i.e CDM, Newmark) has different way to update u, ud, udd from
	 * incremental solution
	 * 
	 * @param U
	 */
	public abstract void updateSolution(IVector U);

	@Override
	public void started(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		super.started(leftHandSideSolver, rightHandSide, u);
		int n = model.getSize();
		/*
		 * initialize system matrix & vectors
		 */
		this.initializeSystemMatricesAndVectors(leftHandSideSolver, rightHandSide, u);

		/*
		 * initialize solution vectors
		 */
		uHat = new ArrayVector(n);
		udHat = new ArrayVector(n);
		uddHat = new ArrayVector(n);

		/*
		 * compute initial condition
		 */
		computeInitialCondition();

		Debugger.warn("Compute initial condition completed");
		Debugger.watch("Time stepping start");
	}

	@Override
	public void timeStepStarted(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		t_old = getCurrentTimeStep();
		t = getNextTimeStep();

		if (getEchoLevel().contain(EchoLevelType.OUTPUT)) {
			Debugger.watch("time = ", getCurrentTimeStep());
		}

		/*
		 * start model time step
		 */
		model.InitializeSolutionStep();
	}

	@Override
	public void iterationStarted(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		/*
		 * start model iteration
		 */
		model.InitalizeNonlinearIteration();
	}

	@Override
	public boolean iterationFinished(Solver leftHandSideSolver, ArrayVector rightHandSide,
			ArrayVector u) {

		leftHandSideSolver.initialize();
		Arrays.fill(rightHandSide.getValues(), 0);
		this.assembleEffectiveSystem(leftHandSideSolver.getA(), rightHandSide);

		if (getEchoLevel().contain(EchoLevelType.AUX2)) {
			Debugger.watch("keff = ", leftHandSideSolver.getA());
			Debugger.watch("reff = ", rightHandSide);
		}

		/*
		 * solve
		 */
		try {
			leftHandSideSolver.solve(rightHandSide.getValues());
		} catch (SolveFailedException e) {
			throw new Error("Solve failed");
		}

		/*
		 * update solution: update uHat, udHat, uddHat
		 */
		this.updateSolution(rightHandSide);
		/*
		 * update model material
		 */
		model.setX(t, uHat.getValues());
		model.FinalizeNonlinearIteration();
		/*
		 * check convergenve criteria
		 */
		double residual = BLAM.norm(BLAM.NORM_TWO, rightHandSide);
		/*
		 * TODO: it's a trivial convergence criteria. Improve this.
		 */
		boolean conv_criteria = (residual < MathUtilities.tol6);

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("error = ", residual);
			Debugger.watch("uHat = ", uHat);
		}

		return conv_criteria;
	}

	@Override
	public void timeStepFinished(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		/*
		 * trigger model to act on material
		 */
		model.FinalizeSolutionStep();
		/*
		 * save solution
		 */
		BLAM.copy(uHat, uHat_old);
		BLAM.copy(udHat, udHat_old);
		BLAM.copy(uddHat, uddHat_old);
	}

}
